#!/usr/bin/python3
# -*- coding: utf-8 -*-
'''
This is the second version of the fusion of outcome of cameras.
In the first version, if results from several cameras appear at the same time, we will fuse them in a fixed weight.
However, in this version, we output them individually, considering of the latency of one of these results in agile environment.
This version may be the final version. 

Write data to a file using multiprocessing causes data fusion.
No matter write in simply or lock file when data is writing using fcntl do not solve the problem.
Therefore, I do not record data using this file.
Instead, I write in sevel files which correspond to the corresponding process.
'''
import serial
import serial.tools.list_ports

import rospy
from geometry_msgs.msg import Transform,Pose,Point32
from cv_bridge import CvBridge, CvBridgeError

import numpy as np
import csv

from bs_img.serial_pack import *

from bs_img.bs_cfg_real_5m import *
 
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import pandas as pd
import traceback

import time
import fcntl


class get_final_res():
    def __init__(self) -> None:
        self.lf_drogue_res_topic = "/bs_debuger/result/lf_drogue"
        self.sf_drogue_res_topic = "/bs_debuger/result/sf_drogue"
        self.wa_drogue_res_topic = "/bs_debuger/result/wa_drogue"
        self.lf_plane_res_topic = "/bs_debuger/result/lf_plane"
        self.sf_plane_res_topic = "/bs_debuger/result/sf_plane"

        self.lf_drogue_res = [np.nan]*3
        self.sf_drogue_res = [np.nan]*3
        self.wa_drogue_res = [np.nan]*3
        self.lf_plane_res = [np.nan]*3
        self.sf_plane_res = [np.nan]*3
        # weight coefficient for the results of lf camera in fusion
        self.coeff = camK_lf[0][0]/(camK_lf[0][0]+camK_sf[0][0])
        self.sleep_interval = 0.0005

        #final result
        self.plane_res = np.array([np.nan, np.nan, np.nan])
        self.drogue_res = np.array([np.nan, np.nan, np.nan])

        self.invalid_frame_num_plane = 0      #compute continus invalid frames
        self.invalid_frame_num_drogue = 0
        self.invalid_frame_max = 5     #if invaild frame < invalid_frame_max, output results of the last frame

        #Kalman Filter
        self.Q = np.eye(3)
        self.R = np.eye(3)
        self.x_plane = np.array([np.nan, np.nan, np.nan])
        self.x_drogue = np.array([np.nan, np.nan, np.nan])
        self.P_plane = np.eye(3)
        self.P_drogue = np.eye(3)

        # csvfile = open ("/home/nuc001/bs_gd_single-master/bag/data.csv","w") 
        # self.writer = csv.writer(csvfile)
        # write2csv(['cam_cls','target_cls','IsValid', 'time','x','y','z'])
        write2csv(['cam_cls','target_cls','IsValid', 'time','x','y','z'])


        self.lf_drogue_res_sub = rospy.Subscriber(self.lf_drogue_res_topic, Transform, self.lf_drogue_res_cb)
        self.sf_drogue_res_sub = rospy.Subscriber(self.sf_drogue_res_topic, Transform, self.sf_drogue_res_cb)
        self.wa_drogue_res_sub = rospy.Subscriber(self.wa_drogue_res_topic, Transform, self.wa_drogue_res_cb)
        self.lf_plane_res_sub = rospy.Subscriber(self.lf_plane_res_topic, Transform, self.lf_plane_res_cb)
        self.sf_plane_res_sub = rospy.Subscriber(self.sf_plane_res_topic, Transform, self.sf_plane_res_cb)


    def lf_drogue_res_cb(self,transform_msg):
        self.lf_drogue_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        lf_drouge_is_visiable = True if not(np.all(np.isnan(self.lf_drogue_res))) else False
        sf_drouge_is_visiable = True if not(np.all(np.isnan(self.sf_drogue_res))) else False
        wa_drouge_is_visiable = True if not(np.all(np.isnan(self.wa_drogue_res))) else False
        if lf_drouge_is_visiable or sf_drouge_is_visiable or wa_drouge_is_visiable:
            self.invalid_frame_num_drogue = 0
            if wa_drouge_is_visiable:
                self.drogue_res = self.wa_drogue_res
            elif lf_drouge_is_visiable and sf_drouge_is_visiable:
                self.drogue_res = np.array(self.lf_drogue_res)*self.coeff + np.array(self.sf_drogue_res)*(1-self.coeff)    
            elif not(lf_drouge_is_visiable) and sf_drouge_is_visiable:
                self.drogue_res = self.sf_drogue_res  
            elif not(sf_drouge_is_visiable) and lf_drouge_is_visiable:
                self.drogue_res = self.lf_drogue_res
            if np.all(np.isnan(self.x_drogue)):   #the first frame
                self.x_drogue = self.drogue_res
            else: 
                z_measure_drogue = np.array(self.drogue_res)
                self.P_drogue = self.P_drogue + self.Q
                K = np.matmul(self.P_drogue, np.linalg.pinv(self.P_drogue + self.R))
                self.x_drogue = self.x_drogue + np.matmul(K, z_measure_drogue - self.x_drogue)
                self.P_drogue = np.matmul(np.eye(3) - K, self.P_drogue)
            write2csv(["lf","drogue",True, time.time(), self.x_drogue[2], self.x_drogue[0], self.x_drogue[1]]) 
        else:
            self.invalid_frame_num_drogue = self.invalid_frame_num_drogue + 1
            if self.invalid_frame_num_drogue <= self.invalid_frame_max:
                self.drogue_res = self.x_drogue
                if not(np.all(np.isnan(self.drogue_res))):
                    write2csv(["lf","drogue",True, time.time(), self.x_drogue[2], self.x_drogue[0], self.x_drogue[1]]) 
                else:
                    write2csv(["lf","drogue",False, time.time(), 0,0,0])
            else:
                write2csv(["lf","drogue",False, time.time(), 0,0,0])



        
    def sf_drogue_res_cb(self,transform_msg):
        self.sf_drogue_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        lf_drouge_is_visiable = True if not(np.all(np.isnan(self.lf_drogue_res))) else False
        sf_drouge_is_visiable = True if not(np.all(np.isnan(self.sf_drogue_res))) else False
        wa_drouge_is_visiable = True if not(np.all(np.isnan(self.wa_drogue_res))) else False
        if lf_drouge_is_visiable or sf_drouge_is_visiable or wa_drouge_is_visiable:
            self.invalid_frame_num_drogue = 0
            if wa_drouge_is_visiable:
                self.drogue_res = self.wa_drogue_res
            elif lf_drouge_is_visiable and sf_drouge_is_visiable:
                self.drogue_res = np.array(self.lf_drogue_res)*self.coeff + np.array(self.sf_drogue_res)*(1-self.coeff)    
            elif not(lf_drouge_is_visiable) and sf_drouge_is_visiable:
                self.drogue_res = self.sf_drogue_res  
            elif not(sf_drouge_is_visiable) and lf_drouge_is_visiable:
                self.drogue_res = self.lf_drogue_res
            if np.all(np.isnan(self.x_drogue)):   #the first frame
                self.x_drogue = self.drogue_res
            else: 
                z_measure_drogue = np.array(self.drogue_res)
                self.P_drogue = self.P_drogue + self.Q
                K = np.matmul(self.P_drogue, np.linalg.pinv(self.P_drogue + self.R))
                self.x_drogue = self.x_drogue + np.matmul(K, z_measure_drogue - self.x_drogue)
                self.P_drogue = np.matmul(np.eye(3) - K, self.P_drogue)
            write2csv(["sf","drogue",True, time.time(), self.x_drogue[2], self.x_drogue[0], self.x_drogue[1]])  
        else:
            self.invalid_frame_num_drogue = self.invalid_frame_num_drogue + 1
            if self.invalid_frame_num_drogue <= self.invalid_frame_max:
                self.drogue_res = self.x_drogue
                if not(np.all(np.isnan(self.drogue_res))):
                    write2csv(["sf","drogue",True, time.time(), self.x_drogue[2], self.x_drogue[0], self.x_drogue[1]])  
                else:
                    write2csv(["sf","drogue",False, time.time(), 0,0,0])  
            else:
                write2csv(["sf","drogue",False, time.time(), 0,0,0])


    def wa_drogue_res_cb(self,transform_msg):
        self.wa_drogue_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        lf_drouge_is_visiable = True if not(np.all(np.isnan(self.lf_drogue_res))) else False
        sf_drouge_is_visiable = True if not(np.all(np.isnan(self.sf_drogue_res))) else False
        wa_drouge_is_visiable = True if not(np.all(np.isnan(self.wa_drogue_res))) else False
        if lf_drouge_is_visiable or sf_drouge_is_visiable or wa_drouge_is_visiable:
            self.invalid_frame_num_drogue = 0
            if wa_drouge_is_visiable:
                self.drogue_res = self.wa_drogue_res
            elif lf_drouge_is_visiable and sf_drouge_is_visiable:
                self.drogue_res = np.array(self.lf_drogue_res)*self.coeff + np.array(self.sf_drogue_res)*(1-self.coeff)    
            elif not(lf_drouge_is_visiable) and sf_drouge_is_visiable:
                self.drogue_res = self.sf_drogue_res  
            elif not(sf_drouge_is_visiable) and lf_drouge_is_visiable:
                self.drogue_res = self.lf_drogue_res
            if np.all(np.isnan(self.x_drogue)):   #the first frame
                self.x_drogue = self.drogue_res
            else: 
                z_measure_drogue = np.array(self.drogue_res)
                self.P_drogue = self.P_drogue + self.Q
                K = np.matmul(self.P_drogue, np.linalg.pinv(self.P_drogue + self.R))
                self.x_drogue = self.x_drogue + np.matmul(K, z_measure_drogue - self.x_drogue)
                self.P_drogue = np.matmul(np.eye(3) - K, self.P_drogue)
            write2csv(["wa","drogue",True, time.time(), self.x_drogue[2], self.x_drogue[0], self.x_drogue[1]])  
        else:
            self.invalid_frame_num_drogue = self.invalid_frame_num_drogue + 1
            if self.invalid_frame_num_drogue <= self.invalid_frame_max:
                self.drogue_res = self.x_drogue
                if not(np.all(np.isnan(self.drogue_res))):
                    write2csv(["wa","drogue",True, time.time(), self.x_drogue[2], self.x_drogue[0], self.x_drogue[1]])     
                else:
                    write2csv(["wa","drogue",False, time.time(), 0,0,0])
            else:
                write2csv(["wa","drogue",False, time.time(), 0,0,0])


    def lf_plane_res_cb(self,transform_msg):
        sf_plane_is_visiable = True if not(np.all(np.isnan(self.sf_plane_res))) else False
        lf_plane_is_visiable = True if not(np.all(np.isnan(self.lf_plane_res))) else False
        if not(sf_plane_is_visiable) and not(lf_plane_is_visiable):
            self.invalid_frame_num_plane = self.invalid_frame_num_plane + 1
            if self.invalid_frame_num_plane <= self.invalid_frame_max:
                self.plane_res = self.x_plane
                if not(np.all(np.isnan(self.plane_res))):
                    write2csv(["lf","plane",True, time.time(),self.plane_res[2],self.plane_res[0],self.plane_res[1]])
                else:
                    write2csv(["lf","plane",False, time.time(),0,0,0])
            else:
                write2csv(["lf","plane",False, time.time(),0,0,0])
        
        else:
            self.invalid_frame_num_plane = 0
            if sf_plane_is_visiable and not(lf_plane_is_visiable):
                self.plane_res = self.sf_plane_res  
            elif lf_plane_is_visiable and not(sf_plane_is_visiable):
                self.plane_res = self.lf_plane_res 
            else:
                self.plane_res = np.array(self.lf_plane_res)*self.coeff + np.array(self.sf_plane_res)*(1-self.coeff)
            if np.all(np.isnan(self.x_plane)):   #the first frame
                self.x_plane = self.plane_res
            else:
                z_measure_plane = np.array(self.plane_res)
                self.P_plane = self.P_plane + self.Q
                K = np.matmul(self.P_plane, np.linalg.pinv(self.P_plane + self.R))
                self.x_plane = self.x_plane + np.matmul(K, z_measure_plane - self.x_plane)
                self.P_plane = np.matmul(np.eye(3) - K, self.P_plane)
            write2csv(["lf","plane",True, time.time(),self.plane_res[2],self.plane_res[0],self.plane_res[1]])


    def sf_plane_res_cb(self,transform_msg):
        self.sf_plane_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        sf_plane_is_visiable = True if not(np.all(np.isnan(self.sf_plane_res))) else False
        lf_plane_is_visiable = True if not(np.all(np.isnan(self.lf_plane_res))) else False
        if not(sf_plane_is_visiable) and not(lf_plane_is_visiable):
            self.invalid_frame_num_plane = self.invalid_frame_num_plane + 1
            if self.invalid_frame_num_plane <= self.invalid_frame_max:
                self.plane_res = self.x_plane
                if not(np.all(np.isnan(self.plane_res))):
                    write2csv(["sf","plane",True, time.time(),self.plane_res[2],self.plane_res[0],self.plane_res[1]])
                else:
                    write2csv(["sf","plane",False, time.time(),0,0,0])
            else:
                write2csv(["sf","plane",False, time.time(),0,0,0])
        
        else:
            self.invalid_frame_num_plane = 0
            if sf_plane_is_visiable and not(lf_plane_is_visiable):
                self.plane_res = self.sf_plane_res  
            elif lf_plane_is_visiable and not(sf_plane_is_visiable):
                self.plane_res = self.lf_plane_res 
            else:
                self.plane_res = np.array(self.lf_plane_res)*self.coeff + np.array(self.sf_plane_res)*(1-self.coeff)
            if np.all(np.isnan(self.x_plane)):   #the first frame
                self.x_plane = self.plane_res
            else:
                z_measure_plane = np.array(self.plane_res)
                self.P_plane = self.P_plane + self.Q
                K = np.matmul(self.P_plane, np.linalg.pinv(self.P_plane + self.R))
                self.x_plane = self.x_plane + np.matmul(K, z_measure_plane - self.x_plane)
                self.P_plane = np.matmul(np.eye(3) - K, self.P_plane)
            write2csv(["sf","plane",True, time.time(),self.plane_res[2],self.plane_res[0],self.plane_res[1]])
        
def write2csv(data):
    with open ("/home/nuc001/bs_gd_single-master/bag/data.csv","a") as file:   # appending
        fcntl.flock(file, fcntl.LOCK_EX)
        try:
            writer = csv.writer(file)
            writer.writerow(data)
        except:
            print("error with opening file")
        finally:
            fcntl.flock(file, fcntl.LOCK_UN)


if __name__ == "__main__":
    rospy.init_node('bs_solver_node', anonymous=True)

    final_res = get_final_res()
    rospy.spin()




